import json

import pandas as pd
from tqdm import tqdm
import numpy as np

with open('唐诗.json', 'r', encoding='utf-8') as f:
    dataset = json.load(f)

voc = dict()

for i, one_data in enumerate(tqdm(dataset)):
    one = one_data[4].replace('\n', '')
    dataset[i] = one
    for word in one:
        voc[word] = 1

dataset = [i for i in dataset if len(i) < 16 * 8 - 1 and len(i) > 8]
voc = ["<|pad|>", "<|aos|>"] + sorted(voc.keys()) + ['<|p_{}|>'.format(i) for i in range(0, 16 * 8 + 1)]
# 多套编码策略 输入是两套编码，输出是一套编码 共三套 编码
voc = pd.DataFrame({"voc": voc,
                    "input_voc1": [0] + np.random.choice(list(range(1, len(voc))), len(voc) - 1,
                                                         replace=False).tolist(),
                    "input_voc2": [0] + np.random.choice(list(range(1, len(voc))), len(voc) - 1,
                                                         replace=False).tolist(),
                    "output_voc": [0] + np.random.choice(list(range(1, len(voc))), len(voc) - 1,
                                                         replace=False).tolist()})
pd.to_pickle(voc, 'voc.pandas_pickle')

for i, one in enumerate(tqdm(dataset)):
    one = ["<|aos|>"] + list(one)
    two = []
    for j in one:
        two.append(j)
        if j in ["，", "。"]:
            two.append("<|p_{}|>".format(len(two)))

    one = two

    if len(one) < 16 * 8 + 1:
        one += ['<|p_{}|>'.format(i) for i in range(len(one), 16 * 8 + 1)]

    else:
        one = one[:16 * 8 + 1]
    one_id0 = []
    one_id1 = []
    one_id2 = []

    for word in one:
        a,b,c = voc.loc[voc["voc"].isin([word]), ["input_voc1","input_voc2","output_voc"]].values[0]

        one_id0.append(a)
        one_id1.append(b)
        one_id2.append(c)

    dataset[i] = [one_id0, one_id1, one_id2]
pd.to_pickle(dataset, 'dataset.pandas_pickle')
